package translationModels;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Random;
import java.util.AbstractMap.SimpleEntry;

import main.Sentence;
import main.TrainingData;

public class IBMModel2 implements IBMModel
{
	private IBMModel1 model1;
	private TrainingData trainingData;
	private HashMap<TParam, Double> tParams;
	private HashMap<QParam, Double> qParams;
	private ArrayList<HashMap<Integer, Integer>> alignments;
	
	public IBMModel2(TrainingData trainingData)
	{
		model1 = new IBMModel1(trainingData);
		this.trainingData = trainingData;
		this.tParams = new HashMap<TParam, Double>();
		this.qParams = new HashMap<QParam, Double>();
		alignments = new ArrayList<HashMap<Integer, Integer>>();
	}
	
	private void reportProgress(int iterations, long done)
	{
		double percent = ((double)done) / (((long)iterations) * trainingData.getSentences().size()) * 100;
		System.out.print("\rModel 2: Completed " + (int)percent + "% of the learning...");
	}
	
	public HashMap<TParam, Double> learnParams(TParamsInitialization initMethod, int iterations)
	{
		tParams.clear();
		Random rnd = new Random();
		
		if (initMethod == TParamsInitialization.Bootstrapping)
			tParams = model1.learnParams(TParamsInitialization.Uniform, iterations);
		
		// Repeat according to the given iterations
		System.out.print("Model 2: Completed 0% of the learning...");
		for (int t = 0; t < iterations; t++)
		{
			HashMap<ArrayList<String>, Double> feCounters = new HashMap<ArrayList<String>, Double>();
			HashMap<String, Double> eCounters = new HashMap<String, Double>();
			HashMap<ArrayList<Integer>, Double> jilmCounters = new HashMap<ArrayList<Integer>, Double>(); 
			HashMap<ArrayList<Integer>, Double> ilmCounters = new HashMap<ArrayList<Integer>, Double>();
			int sentencesCounter = 0;
			
			for (SimpleEntry<Sentence, Sentence> entry : trainingData.getSentences())
			{
				sentencesCounter++;
				
				for (int i = 0; i < entry.getKey().length(); i++)
				{
					String f = entry.getKey().getWord(i);
					for (int j = -1; j < entry.getValue().length(); j++)
					{
						String e = entry.getValue().getWord(j);
						TParam tParam = new TParam(f, e);
						QParam qParam = new QParam(j,i,entry.getValue().length(),entry.getKey().length());
						
						// Calculate the delta
						double sum = 0;
						for (int k = -1; k < entry.getValue().length(); k++)
						{
							QParam currentQ = new QParam(k,i,entry.getValue().length(),entry.getKey().length());
							TParam current = new TParam(f, entry.getValue().getWord(k));
							if (!tParams.containsKey(current))
							{
								if (initMethod == TParamsInitialization.Random)
									tParams.put(current, rnd.nextDouble());
								else
									tParams.put(current, 1.0);
							}
							if (!qParams.containsKey(currentQ))
							{
								sum += tParams.get(current);
							}
							else
							sum += qParams.get(currentQ) * tParams.get(current);
						}
						double delta;
						if (!qParams.containsKey(qParam))
							delta = tParams.get(tParam) / sum;
						else
							delta = (qParams.get(qParam) * tParams.get(tParam)) / sum;
						
						// Update the counters
						ArrayList<String> list = tParam.toArrayList();
						if (!feCounters.containsKey(list))
							feCounters.put(list, 0.0);
						feCounters.put(list, feCounters.get(list) + delta);
						if (!eCounters.containsKey(e))
							eCounters.put(e, 0.0);
						eCounters.put(e, eCounters.get(e) + delta);
						ArrayList<Integer> list2 = qParam.toArrayList();
						if(!jilmCounters.containsKey(list2))
							jilmCounters.put(list2, 0.0);
						jilmCounters.put(list2, jilmCounters.get(list2) + delta);
						ArrayList<Integer> sublist = qParam.getILMList();
						if(!ilmCounters.containsKey(sublist))
							ilmCounters.put(sublist, 0.0);
						ilmCounters.put(sublist, ilmCounters.get(sublist) + delta);
					}
				}
				
				reportProgress(iterations, (long)t * trainingData.getSentences().size() + sentencesCounter);
			}
			
			// Update t from counters
			for (TParam tParam : tParams.keySet())
				tParams.put(tParam, feCounters.get(tParam.toArrayList()) / eCounters.get(tParam.e));
			
			// Update q from counters
			for (QParam qParam : qParams.keySet())
				qParams.put(qParam, jilmCounters.get(qParam.toArrayList()) / ilmCounters.get(qParam.getILMList()));
		}
		System.out.println("\rModel 2: Completed 100% of the learning...");
		
		return tParams;
	}
	
	public ArrayList<HashMap<Integer, Integer>> calculateAlignments()
	{
		for (SimpleEntry<Sentence, Sentence> entry : trainingData.getSentences())
		{
			Sentence f = entry.getKey();
			Sentence e = entry.getValue();
			int m = f.length();
			int l = e.length();
			alignments.add(new HashMap<Integer, Integer>());
			HashMap<Integer, Integer> currentAlignment = alignments.get(alignments.size() - 1);
			for (int i = 0; i < f.length(); i++)
			{	
				int best = -1;
				for (int j = 0; j < e.length(); j++)
				{
					double currScore = tParams.get(new TParam(f.getWord(i), e.getWord(j)));
					if(qParams.get(new QParam(j, i, l, m)) != null)
							currScore *= qParams.get(new QParam(j, i, l, m));
					
					double bestScore = tParams.get(new TParam(f.getWord(i), e.getWord(best)));
					if(qParams.get(new QParam(best,i,l,m)) != null)
					{
						bestScore *= qParams.get(new QParam(best,i,l,m));
					}
					if(currScore > bestScore)
						best = j;
				}
				currentAlignment.put(i, best);
			}
		}
		
		return alignments;
	}
}
